"""
收益率分析器模块
"""
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from loguru import logger
from .analyzer import Analyzer


class ReturnsAnalyzer(Analyzer):
    """
    收益率分析器，分析回测结果的收益情况
    """
    
    def __init__(self, trades, equity_curve):
        """
        初始化收益率分析器
        
        Parameters
        ----------
        trades : pandas.DataFrame
            交易记录
        equity_curve : pandas.DataFrame
            资金曲线
        """
        super().__init__(trades, equity_curve)
        self._returns = None
        self._cum_returns = None
        self._drawdowns = None
        self._stats = None
    
    def analyze(self):
        """
        分析回测结果
        
        Returns
        -------
        dict
            分析结果
        """
        # 计算收益率
        self._returns = self._equity_curve['equity'].pct_change().fillna(0)
        
        # 计算累计收益率
        self._cum_returns = (1 + self._returns).cumprod() - 1
        
        # 计算回撤
        self._drawdowns = self._calculate_drawdowns()
        
        # 计算统计指标
        self._stats = self._calculate_stats()
        
        return self._stats
    
    def _calculate_drawdowns(self):
        """
        计算回撤
        
        Returns
        -------
        pandas.Series
            回撤序列
        """
        # 计算历史最高点
        cumulative = (1 + self._returns).cumprod()
        running_max = cumulative.cummax()
        
        # 计算回撤
        drawdowns = (cumulative / running_max) - 1
        
        return drawdowns
    
    def _calculate_stats(self):
        """
        计算统计指标
        
        Returns
        -------
        dict
            统计指标
        """
        # 确保有足够的数据
        if len(self._returns) < 2:
            return {
                'total_return': 0,
                'annual_return': 0,
                'sharpe_ratio': 0,
                'max_drawdown': 0,
                'volatility': 0,
                'win_rate': 0,
                'profit_factor': 0
            }
        
        # 计算总收益率
        total_return = self._cum_returns.iloc[-1]
        
        # 计算年化收益率（假设252个交易日）
        days = (self._equity_curve.index[-1] - self._equity_curve.index[0]).days
        if days <= 0:
            annual_return = 0
        else:
            annual_return = (1 + total_return) ** (252 / days) - 1
        
        # 计算波动率
        volatility = self._returns.std() * np.sqrt(252)
        
        # 计算夏普比率（假设无风险利率为0）
        if volatility > 0:
            sharpe_ratio = annual_return / volatility
        else:
            sharpe_ratio = 0
        
        # 计算最大回撤
        max_drawdown = self._drawdowns.min()
        
        # 计算胜率
        win_days = (self._returns > 0).sum()
        total_days = len(self._returns)
        win_rate = win_days / total_days if total_days > 0 else 0
        
        # 计算盈亏比
        gains = self._returns[self._returns > 0].sum()
        losses = -self._returns[self._returns < 0].sum()
        profit_factor = gains / losses if losses > 0 else float('inf')
        
        return {
            'total_return': total_return,
            'annual_return': annual_return,
            'sharpe_ratio': sharpe_ratio,
            'max_drawdown': max_drawdown,
            'volatility': volatility,
            'win_rate': win_rate,
            'profit_factor': profit_factor
        }
    
    def _get_stats_text(self):
        """
        获取统计指标文本
        
        Returns
        -------
        str
            统计指标文本
        """
        if self._stats is None:
            return ""
            
        return (
            f"总收益率: {self._stats['total_return']:.2%}\n"
            f"年化收益率: {self._stats['annual_return']:.2%}\n"
            f"夏普比率: {self._stats['sharpe_ratio']:.2f}\n"
            f"最大回撤: {self._stats['max_drawdown']:.2%}\n"
            f"波动率: {self._stats['volatility']:.2%}\n"
            f"胜率: {self._stats['win_rate']:.2%}\n"
            f"盈亏比: {self._stats['profit_factor']:.2f}"
        )
